#! /bin/bash

# This is a unit test program for kpartx, in particular for deleting partitions.
#
# The rationale is the following:
#
#  1) kpartx should delete all mappings it created beforehand.
#  2) kpartx should handle partitions on dm devices and other devices
#     (e.g. loop devices) equally well.
#  3) kpartx should only delete "partitions", which are single-target
#     linear mappings into a block device. Other maps should not be touched.
#  4) kpartx should only delete mappings it created itself beforehand.
#     In particular, it shouldn't delete LVM LVs, even if they are fully
#     contained in the block device at hand and thus look like partitions
#     in the first place. (For historical compatibility reasons, we allow
#     such mappings to be deleted with the -f/--force flag).
#  5) DM map names may be changed, thus kpartx shouldn't rely on them to
#     check whether a mapping is a partition of a particular device. It is
#     legal for a partition of /dev/loop0 to be named "loop0".

# Note: This program tries hard to clean up, but if tests fail,
# stale DM or loop devices may keep lurking around.

# Set WORKDIR in environment to existing dir to for persistence
# WARNING:  existing files will be truncated.
# If empty, test will be done in temporary dir
: ${WORKDIR:=}
# Set this environment variable to test an alternative kpartx executable
: ${KPARTX:=}
# Options to pass to kpartx always
: ${KPARTX_OPTS:=-s}
# Time to wait for device nodes to appear (microseconds)
# Waiting is only needed if "s" is not in $KPARTX_OPTS
: ${WAIT_US:=0}

# IMPORTANT: The ERR trap is essential for this program to work correctly!
trap 'LINE=$LINENO; trap - ERR; echo "== error in $BASH_COMMAND on line $LINE ==" >&2; exit 1' ERR
trap 'cleanup' 0

CLEANUP=:
cleanup() {
    trap - ERR
    trap - 0
    if [[ $OK ]]; then
	echo == all tests completed successfully == >&2
    else
	echo == step $STEP failed == >&2
    fi
    eval "$CLEANUP" &>/dev/null
}

push_cleanup() {
    CLEANUP="$@;$CLEANUP"
}

pop_cleanup() {
    # CAUTION: simplistic
    CLEANUP=${CLEANUP#*;}
}

step() {
    STEP="$@"
    echo == Test step: $STEP == >&2
}

mk_partitions() {
    parted -s $1 mklabel msdos
    parted -s -- $1 mkpart prim ext2 1MiB -1s
}

wipe_ptable() {
    dd if=/dev/zero of=$1 bs=1b count=1
}

step preparation

[[ $UID -eq 0 ]]
[[ $KPARTX ]] || {
    if [[ -x $PWD/kpartx/kpartx ]]; then
	KPARTX=$PWD/kpartx/kpartx
    else
	KPARTX=$(which kpartx)
    fi
}
[[ $KPARTX ]]

FILE1=kpartx1
FILE2=kpartx2
FILE3=kpartx3
FILE4=kpartx4
SIZE=$((1024*1024*1024))  # use bytes as units here
SECTSIZ=512
OFFS=32                # offset of linear mapping into dev, sectors
VG=kpvg  # volume group name
LV=kplv  # logical vol name
LVMCONF='devices { filter = [ "a|/dev/loop.*|", r".*" ] }'

OK=

[[ $WORKDIR ]] || {
    WORKDIR=$(mktemp -d /tmp/kpartx-XXXXXX)
    push_cleanup 'rm -rf $WORKDIR'
}

push_cleanup "cd $PWD"
cd "$WORKDIR"

step "create loop devices"
truncate -s $SIZE $FILE1
truncate -s $SIZE $FILE2
truncate -s $SIZE $FILE3
truncate -s $SIZE $FILE4

LO1=$(losetup -f $FILE1 --show)
push_cleanup 'losetup -d $LO1'
LO2=$(losetup -f $FILE2 --show)
push_cleanup 'losetup -d $LO2'
LO3=$(losetup -f $FILE3 --show)
push_cleanup 'losetup -d $LO3'
LO4=$(losetup -f $FILE4 --show)
push_cleanup 'losetup -d $LO4'

[[ $LO1 && $LO2 && $LO3 && $LO4 && -b $LO1 && -b $LO2 && -b $LO3 && -b $LO4 ]]
DEV1=$(stat -c "%t:%T" $LO1)
DEV2=$(stat -c "%t:%T" $LO2)
DEV3=$(stat -c "%t:%T" $LO3)

usleep $WAIT_US

step "create DM devices (spans)"
# Create two linear mappings spanning two loopdevs.
# One of them gets a pathological name colliding with
# the loop device name.
# These mappings must not be removed by kpartx.
# They also serve as DM devices to test partition removal on those.

TABLE="\
0 $((SIZE/SECTSIZ-OFFS)) linear $DEV1 $OFFS
$((SIZE/SECTSIZ-OFFS)) $((SIZE/SECTSIZ-OFFS)) linear $DEV2 $OFFS"

SPAN1=kpt
SPAN2=$(basename $LO2)
dmsetup create $SPAN1 <<<"$TABLE"
push_cleanup 'dmsetup remove -f $SPAN1'

dmsetup create $SPAN2 <<<"$TABLE"
push_cleanup 'dmsetup remove -f $SPAN2'

# This is a non-kpartx pseudo "partition" mapping
USER1=user1
push_cleanup 'dmsetup remove -f $USER1'
dmsetup create $USER1 <<EOF
0 $((SIZE/SECTSIZ-OFFS)) linear $DEV1 $OFFS
EOF

usleep $WAIT_US
[[ -b /dev/mapper/$SPAN1 ]]
[[ -b /dev/mapper/$SPAN2 ]]
[[ -b /dev/mapper/$USER1 ]]

step "create vg on $LO3"
# On the 3rd loop device, we create a VG and an LV
# The LV should not be removed by kpartx.
pvcreate --config "$LVMCONF" -f $LO3
vgcreate --config "$LVMCONF" $VG $LO3
push_cleanup 'vgremove --config "$LVMCONF" -f $VG'
lvcreate --config "$LVMCONF" -L $((SIZE/2))B -n $LV $VG
push_cleanup 'lvremove --config "$LVMCONF" -f $VG/$LV'
usleep $WAIT_US

[[ -b /dev/mapper/$VG-$LV ]]

# dmsetup table /dev/mapper/$VG-$LV
# dmsetup info /dev/mapper/$VG-$LV

step "create partitions on loop devices"

mk_partitions $LO1
mk_partitions $LO2
mk_partitions $LO4

# Test invocation of kpartx with regular file here
LO2P1=/dev/mapper/$(basename $LO2)-foo1
$KPARTX $KPARTX_OPTS -a -p -foo $FILE2
[[ -b $LO2P1 ]]
push_cleanup 'dmsetup remove -f $(basename $LO2P1)'

step "remove partitions with deleted ptable"
wipe_ptable $LO2
$KPARTX $KPARTX_OPTS -d $LO2
[[ ! -b $LO2P1 ]]

mk_partitions $LO2
$KPARTX $KPARTX_OPTS -a -p -foo $FILE2
[[ -b $LO2P1 ]]

LO1P1=/dev/mapper/$(basename $LO1)-eggs1
$KPARTX $KPARTX_OPTS -a -p -eggs $LO1
push_cleanup 'dmsetup remove -f $(basename $LO1P1)'

usleep $WAIT_US
[[ -b $LO1P1 ]]
[[ -b $LO2P1 ]]

# dmsetup info $LO2P1

# Set pathological name for partition on $LO1 (same as loop device itself)
dmsetup rename $(basename $LO1P1) $(basename $LO1)
LO1P1=/dev/mapper/$(basename $LO1)
pop_cleanup
push_cleanup 'dmsetup remove -f $(basename $LO1P1)'

# dmsetup info $LO1P1

step "create partitions on DM devices"
mk_partitions /dev/mapper/$SPAN2

$KPARTX $KPARTX_OPTS -a -p -bar /dev/mapper/$SPAN2
SPAN2P1=/dev/mapper/${SPAN2}-bar1

# udev rules may have created partition mappings without UUIDs
# which aren't removed by default (if system standard kpartx doesn't
# set the UUID). Remove them using -f
push_cleanup '$KPARTX $KPARTX_OPTS -f -d /dev/mapper/$SPAN2'
push_cleanup 'dmsetup remove -f $(basename $SPAN2P1)'

$KPARTX $KPARTX_OPTS -a -p -spam /dev/mapper/$SPAN1
SPAN1P1=/dev/mapper/${SPAN1}-spam1
# see above
push_cleanup '$KPARTX $KPARTX_OPTS -f -d /dev/mapper/$SPAN1'
push_cleanup 'dmsetup remove -f $(basename $SPAN1P1)'

usleep $WAIT_US
[[ -b $SPAN2P1 ]]
[[ -b $SPAN1P1 ]]

step "rename partitions on DM device to default"
$KPARTX $KPARTX_OPTS -u /dev/mapper/$SPAN1
[[ ! -b ${SPAN1P1} ]]
# This assumes that $SPAN1 ends in a non-digit
[[ -b ${SPAN1P1//-spam/} ]]

step "rename partitions on DM device back from default"
$KPARTX $KPARTX_OPTS -u -p -spam /dev/mapper/$SPAN1
[[ -b ${SPAN1P1} ]]
[[ ! -b ${SPANP1//-foo/} ]]

step "delete partitions on DM devices"
$KPARTX $KPARTX_OPTS -d /dev/mapper/$SPAN1 >&2
usleep $WAIT_US

[[ -b $SPAN2P1 ]]
[[ -b $LO1P1 ]]
[[ -b $LO2P1 ]]
[[ ! -b $SPAN1P1 ]]

$KPARTX $KPARTX_OPTS -d /dev/mapper/$SPAN2
usleep $WAIT_US

[[ -b $LO1P1 ]]
[[ -b $LO2P1 ]]
[[ ! -b $SPAN2P1 ]]

step "rename partitions on loop device"
$KPARTX $KPARTX_OPTS -u -p -spam $LO2
[[ ! -b ${LO2P1} ]]
[[ -b ${LO2P1//-foo/-spam} ]]

step "rename partitions on loop device back"
$KPARTX $KPARTX_OPTS -u -p -foo $LO2
[[ -b ${LO2P1} ]]
[[ ! -b ${LO2P1//-foo/-spam} ]]

step "rename partitions on loop device to default"
$KPARTX $KPARTX_OPTS -u $LO2
#read a
[[ ! -b ${LO2P1} ]]
# $LO1 ends in a digit
[[ -b ${LO2P1//-foo/p} ]]

step "rename partitions on loop device back from default"
$KPARTX $KPARTX_OPTS -u -p -foo $LO2
[[ -b ${LO2P1} ]]
[[ ! -b ${LO2P1//-foo/p} ]]

step "rename partitions on loop devices"
$KPARTX $KPARTX_OPTS -u -p spam $LO2

step "delete partitions on loop devices"

$KPARTX $KPARTX_OPTS -d $LO3

# This will also delete the loop device
$KPARTX $KPARTX_OPTS -d $FILE2
$KPARTX $KPARTX_OPTS -d $LO1
usleep $WAIT_US

# ls -l /dev/mapper
[[ ! -b $LO1P1 ]]
pop_cleanup
[[ ! -b $LO2P1 ]]
pop_cleanup
# spans should not have been removed
[[ -b /dev/mapper/$SPAN1 ]]
[[ -b /dev/mapper/$SPAN2 ]]
[[ -b /dev/mapper/$USER1 ]]
# LVs neither
[[ -b /dev/mapper/$VG-$LV ]]

step "delete partitions on $LO3 with -f"

$KPARTX $KPARTX_OPTS -f -d $LO3
# -d -f should delete the LV, too
[[ ! -b /dev/mapper/$VG-$LV ]]
[[ -b /dev/mapper/$SPAN1 ]]
[[ -b /dev/mapper/$SPAN2 ]]

step "test kpartx creation/deletion on an image file with no existing loopdev"
losetup -d $LO4

OUTPUT=$($KPARTX $KPARTX_OPTS -v -a $FILE4 2>&1)
read loop dm < \
     <(sed -n  's/^add map \(loop[0-9]*\)p1 ([0-9]*:\([0-9]*\)).*$/\1 dm-\2/p' \
	   <<<$OUTPUT)
[[ $dm && $loop ]]
push_cleanup "dmsetup remove -f /dev/$dm"
push_cleanup "losetup -d /dev/$loop"

[[ -b /dev/mapper/${loop}p1 ]]
$KPARTX -d $KPARTX_OPTS $FILE4
[[ ! -b /dev/mapper/${loop}p1 ]]
# /dev/$loop is _not_ automatically deleted
[[ -b /dev/${loop} ]]

OK=yes
